import os
import sys
import argparse
from pathlib import Path
import shutil


# Add project root to Python path
project_root = str(Path(__file__).parent.parent.parent)
if project_root not in sys.path:
    sys.path.insert(0, project_root)

from devdeploy.inference.inference import create_inference
from PIL import Image

# 解析命令行参数
def parse_args():
    parser = argparse.ArgumentParser(description='图像分类推理测试脚本')
    parser.add_argument('--checkpoint', type=str, default='/Users/zhanghaining/git/SteelAiLab/out_models/resnet18_zhn/fullmodel_iter_best.onnx',
                        help='模型检查点路径 (默认: out_models/resnet18/fullmodel_final.pth)')
    parser.add_argument('--img_dir', type=str, default='/Users/zhanghaining/git/SteelAiLab/data/classify/train/50/',
                        help='测试图片目录路径 (默认: data/classify/val/25)')
    return parser.parse_args()

# 获取命令行参数
args = parse_args()
CHECKPOINT_PATH = args.checkpoint
IMG_DIR = args.img_dir

# 创建推理器实例
infer = create_inference(CHECKPOINT_PATH, device='cpu')

def test_dir():
    """测试目录批量推理功能"""
    print("\n" + "="*80)
    print("目录批量推理测试")
    print("="*80)
    
    # 批量推理
    results = infer.infer_batch(IMG_DIR)

    # 打印结果
    for img_name, result in results.items():
        print(f'class_id: {result["class_id"]}  class_name: {result["class_name"]} 置信度: {result["confidence"]:.3f}      {img_name}')

# 单张图片推理测试
def test_single_image_inference():
    """测试单张图片推理功能"""
    
    # 获取测试图片列表
    image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif')
    test_images = [
        f for f in os.listdir(IMG_DIR) 
        if f.lower().endswith(image_extensions)
    ]
    
    if not test_images:
        print(f"错误：在目录 {IMG_DIR} 中没有找到测试图片")
        return
    
    print(f"找到 {len(test_images)} 张测试图片")
    print("-" * 60)
     
    # 测试所有图片的单张推理
    for i, img_name in enumerate(test_images):
        img_path = os.path.join(IMG_DIR, img_name)
        
        try:
            # 单张图片推理
            result = infer.infer_single(img_path, return_prob=True)
            
            # 打印结果
            print(f'class_id: {result["class_id"]}  class_name: {result["class_name"]} 置信度: {result["confidence"]:.3f}      {img_name}')
            
        except Exception as e:
            print(f"推理失败 {img_name}: {str(e)}")

def test_model_info():
    """测试模型信息获取功能"""
    print(f"\n模型信息:")
    print("-" * 40)
    
    try:
        # 获取任务类型
        task_type = infer.get_task_type()
        print(f"任务类型: {task_type}")
        
        # 获取模型信息
        model_info = infer.get_model_info()
        print(f"模型信息: {model_info}")
        
        # 验证模型
        is_valid = infer.validate_model()
        print(f"模型验证: {'通过' if is_valid else '失败'}")
        
    except Exception as e:
        print(f"获取模型信息失败: {str(e)}")

def test_pil_image_input():
    """测试使用PIL Image对象作为输入的推理功能"""
    print("\n" + "="*80)
    print("PIL Image输入测试")
    print("="*80)
    
    # 获取所有测试图片
    image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif')
    test_images = [
        f for f in os.listdir(IMG_DIR) 
        if f.lower().endswith(image_extensions)
    ]
    
    if not test_images:
        print(f"错误：在目录 {IMG_DIR} 中没有找到测试图片")
        return
    
    print(f"找到 {len(test_images)} 张测试图片")
    print("-" * 80)
    
    # 对每张图片进行测试
    for i, test_image in enumerate(test_images, 1):
        img_path = os.path.join(IMG_DIR, test_image)
        
        try:
            # 打开图片
            img = Image.open(img_path)
            
            # 使用PIL Image对象进行推理
            result = infer.infer_single(img, return_prob=True)
            
            # 打印结果
            print(f'class_id: {result["class_id"]}  class_name: {result["class_name"]} 置信度: {result["confidence"]:.3f}      {test_image}')
            
        except Exception as e:
            print(f"测试失败 {test_image}: {str(e)}")

def test_inference_and_draw():
    """测试带绘图的单张图片推理功能，并将结果保存在 devdeploy/temp/ 文件夹中"""
    print("\n" + "="*80)
    print("带绘图的单张图片推理测试")
    print("="*80)

    # 创建保存结果的目录
    output_dir = "devdeploy/temp"
    # 如果目录已存在，先删除
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    os.makedirs(output_dir)
    print(f"推理结果将保存在 {output_dir} 目录中")

    # 获取测试图片列表
    image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif')
    test_images = [
        f for f in os.listdir(IMG_DIR)
        if f.lower().endswith(image_extensions)
    ]

    if not test_images:
        print(f"错误：在目录 {IMG_DIR} 中没有找到测试图片")
        return

    print(f"找到 {len(test_images)} 张测试图片进行绘图")
    print("-" * 80)

    # 对每张图片进行测试
    for img_name in test_images:
        img_path = os.path.join(IMG_DIR, img_name)
        try:
            # 调用带绘图功能的推理
            result = infer.infer_single(img_path, draw_result=True)

            # 检查是否返回了绘图结果
            if 'draw' in result and isinstance(result['draw'], Image.Image):
                # 保存带标注的图片
                output_path = os.path.join(output_dir, f"result_{img_name}")
                result['draw'].save(output_path)
                
                # 打印结果
                print(f'class_id: {result["class_id"]}  class_name: {result["class_name"]} 置信度: {result["confidence"]:.3f} -> {output_path}')
            else:
                print(f"推理成功，但未返回绘图: {img_name}")

        except Exception as e:
            print(f"处理图片失败 {img_name}: {str(e)}")

# 执行测试
if __name__ == "__main__":
    # 测试目录批量推理
    # test_dir()
    
    # 测试单张图片推理
    test_single_image_inference()
    
    # 测试PIL Image输入
    # test_pil_image_input()
    
    # 测试模型信息
    # test_model_info()

    # 测试带绘图的单张图片推理
    # test_inference_and_draw()

    # print("\n" + "="*80)
    # print("所有测试完成")
    # print("="*80)

